{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Title: #Distribute Candies Among Children III"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Difficulty: #Hard"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Category Title: #Algorithms"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Tag Slug: #math #combinatorics"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Name Translated: #数学 #组合数学"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Solution Name: distributeCandies"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Translated Title: #给小朋友们分糖果 III"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Translated Content:\n",
    "<p>你被给定两个正整数 <code>n</code> 和 <code>limit</code>。</p>\n",
    "\n",
    "<p>返回 <em>在每个孩子得到不超过&nbsp;</em><code>limit</code><em> 个糖果的情况下，将</em> <code>n</code> <em>个糖果分发给</em>&nbsp;<code>3</code> <em>个孩子的&nbsp;<strong>总方法数</strong>。</em></p>\n",
    "\n",
    "<p>&nbsp;</p>\n",
    "\n",
    "<p><b>示例 1:</b></p>\n",
    "\n",
    "<pre>\n",
    "<b>输入：</b>n = 5, limit = 2\n",
    "<b>输出：</b>3\n",
    "<b>解释：</b>有 3 种方式将 5 个糖果分发给 3 个孩子，使得每个孩子得到不超过 2 个糖果：(1, 2, 2), (2, 1, 2) 和 (2, 2, 1)。\n",
    "</pre>\n",
    "\n",
    "<p><b>示例 2:</b></p>\n",
    "\n",
    "<pre>\n",
    "<b>输入：</b>n = 3, limit = 3\n",
    "<b>输出：</b>10\n",
    "<b>解释：</b>有 10 种方式将 3 个糖果分发给 3 个孩子，使得每个孩子得到不超过 3 个糖果：(0, 0, 3), (0, 1, 2), (0, 2, 1), (0, 3, 0), (1, 0, 2), (1, 1, 1), (1, 2, 0), (2, 0, 1), (2, 1, 0) 和 (3, 0, 0)。\n",
    "</pre>\n",
    "\n",
    "<p>&nbsp;</p>\n",
    "\n",
    "<p><b>提示：</b></p>\n",
    "\n",
    "<ul>\n",
    "\t<li><code>1 &lt;= n &lt;= 10<sup>8</sup></code></li>\n",
    "\t<li><code>1 &lt;= limit &lt;= 10<sup>8</sup></code></li>\n",
    "</ul>\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Description: [distribute-candies-among-children-iii](https://leetcode.cn/problems/distribute-candies-among-children-iii/description/)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Solutions: [distribute-candies-among-children-iii](https://leetcode.cn/problems/distribute-candies-among-children-iii/solutions/)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "test_cases = ['5\\n2', '3\\n3']"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def distributeCandies(self, n: int, limit: int) -> int:\n",
    "        if limit > n:\n",
    "            return math.comb(n+2, 2)\n",
    "        \n",
    "        @cache\n",
    "        def dp(i, k):\n",
    "            if k == 0:\n",
    "                return min(i, limit) - max(0, i - limit) + 1\n",
    "            \n",
    "            ans = 0\n",
    "            for j in reversed(range(limit+1)):\n",
    "                if i - j > 2 * limit: break\n",
    "                ans += dp(i - j, k - 1)\n",
    "            return ans\n",
    "        \n",
    "        return dp(n, 1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def distributeCandies(self, n: int, limit: int) -> int:\n",
    "        ans = 0\n",
    "        for i in range(n+1):\n",
    "            for j in range(n+1):\n",
    "                k = n - i - j\n",
    "                if k >= 0 and i <= limit and j <= limit and k <= limit:\n",
    "                    ans += 1\n",
    "        return ans"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def distributeCandies(self, n: int, limit: int) -> int:\n",
    "        ans = 0\n",
    "        for i in range(limit+1):\n",
    "            for j in range(limit+1):\n",
    "                for k in range(limit+1):\n",
    "                    if i+j+k == n:\n",
    "                        ans += 1\n",
    "        return ans"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\r\n",
    "    def distributeCandies(self, n: int, limit: int) -> int:\r\n",
    "        return self.solve2(n, limit)\r\n",
    "    \"\"\"\r\n",
    "        1. Brute Force\r\n",
    "    \"\"\"\r\n",
    "    def solve1(self, n: int, limit: int) -> int:\r\n",
    "        ans = 0\r\n",
    "        for i in range(limit+1):\r\n",
    "            for j in range(limit+1):\r\n",
    "                k = n - i - j\r\n",
    "                if k >= 0 and k <= limit:\r\n",
    "                    ans += 1\r\n",
    "        return ans\r\n",
    "    \"\"\"\r\n",
    "        2. 排容原理\r\n",
    "        a. 所有方案數：\r\n",
    "            H(3, n) = C(n+2, n) = C(n+2, 2)\r\n",
    "        b. 至少一人拿到超過 limit 個糖果的方案數：\r\n",
    "            (先分(limit+1)給他，剩下的分給所有人)\r\n",
    "            3 * H(3, n-(limit+1)) = 3 * C(n−(limit+1)+2, 2)\r\n",
    "        c. 至少兩人拿到超過 limit 個糖果的方案數：\r\n",
    "            3 * C(n−2⋅(limit+1)+2, 2)\r\n",
    "    \"\"\"\r\n",
    "\r\n",
    "    def solve2(self, n: int, limit: int) -> int:\r\n",
    "        if n > 3 * limit: # 三人都拿到超過 limit 個糖果，沒有方案\r\n",
    "            return 0\r\n",
    "        ans = comb(n + 2, 2)\r\n",
    "        if n >= (limit + 1): # 至少一人拿到超過 limit 個糖果\r\n",
    "            ans -= 3 * comb(n - limit + 1, 2)\r\n",
    "        if n >= 2 * (limit + 1): # 至少兩人拿到超過 limit 個糖果\r\n",
    "            ans += 3 * comb(n - 2 * limit, 2)\r\n",
    "        return ans"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def distributeCandies(self, n: int, limit: int) -> int:\n",
    "        @cache\n",
    "        def f(i,n):\n",
    "            if i==2 or n<0:\n",
    "                if n>=0 and n<=limit:\n",
    "                    return 1\n",
    "                return 0\n",
    "            m=0\n",
    "            for j in range(limit+1):\n",
    "                m+=f(i+1,n-j)\n",
    "            return m\n",
    "        return f(0,n)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def distributeCandies(self, n: int, limit: int) -> int:\n",
    "        \n",
    "        ans = 0\n",
    "        for i in range(0, min(n, limit)+1):\n",
    "            ans += max(0, min(limit, n - i) - max(0, n - i - limit) + 1)\n",
    "        \n",
    "        return ans"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "def c2(n):\n",
    "    return n * (n - 1) // 2 if n > 1 else 0\n",
    "\n",
    "class Solution:\n",
    "    def distributeCandies(self, n: int, limit: int) -> int:\n",
    "        return c2(n + 2) - 3 * c2(n - limit + 1) + 3 * c2(n - 2 * limit) - c2(n - 3 * limit - 1)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\r\n",
    "    def distributeCandies(self, n: int, limit: int) -> int:\r\n",
    "        return self.solve2(n, limit)\r\n",
    "    \"\"\"\r\n",
    "        1. Brute Force\r\n",
    "    \"\"\"\r\n",
    "    def solve1(self, n: int, limit: int) -> int:\r\n",
    "        ans = 0\r\n",
    "        for i in range(limit+1):\r\n",
    "            for j in range(limit+1):\r\n",
    "                k = n - i - j\r\n",
    "                if k >= 0 and k <= limit:\r\n",
    "                    ans += 1\r\n",
    "        return ans\r\n",
    "    \"\"\"\r\n",
    "        2. 排容原理\r\n",
    "        a. 所有方案數：\r\n",
    "            H(3, n) = C(n+2, n) = C(n+2, 2)\r\n",
    "        b. 至少一人拿到超過 limit 個糖果的方案數：\r\n",
    "            (先分(limit+1)給他，剩下的分給所有人)\r\n",
    "            3 * H(3, n-(limit+1)) = 3 * C(n−(limit+1)+2, 2)\r\n",
    "        c. 至少兩人拿到超過 limit 個糖果的方案數：\r\n",
    "            3 * C(n−2⋅(limit+1)+2, 2)\r\n",
    "    \"\"\"\r\n",
    "\r\n",
    "    def solve2(self, n: int, limit: int) -> int:\r\n",
    "        if n > 3 * limit: # 三人都拿到超過 limit 個糖果，沒有方案\r\n",
    "            return 0\r\n",
    "        ans = comb(n + 2, 2)\r\n",
    "        if n >= (limit + 1): # 至少一人拿到超過 limit 個糖果\r\n",
    "            ans -= 3 * comb(n - limit + 1, 2)\r\n",
    "        if n >= 2 * (limit + 1): # 至少兩人拿到超過 limit 個糖果\r\n",
    "            ans += 3 * comb(n - 2 * limit, 2)\r\n",
    "        return ans"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def distributeCandies(self, n: int, limit: int) -> int:\n",
    "        dp = list()\n",
    "        for i in range(min(n + 1, limit + 1)):\n",
    "            dp.append(i + 1)\n",
    "        x = limit + 1\n",
    "        while n >= x and x <= 2 * limit:\n",
    "            dp.append(2 * limit - x + 1)\n",
    "            x += 1\n",
    "        ans = 0\n",
    "        # print(dp)\n",
    "        for i in range(min(n + 1, limit + 1)):\n",
    "            if n <= 2 * limit + i:\n",
    "                ans += dp[n - i]\n",
    "        return ans \n",
    "        "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def distributeCandies(self, n: int, limit: int) -> int:\n",
    "        if 3 * limit < n:\n",
    "            return 0\n",
    "        cnt = [0] * (n + 1)\n",
    "        for i in range(min(n, 2 * limit) + 1):\n",
    "            if i > limit:\n",
    "                cnt[i] = 2 * limit - i + 1\n",
    "            else:\n",
    "                cnt[i] = i + 1\n",
    "        return sum(cnt[i] for i in range(n + 1) if n - i <= limit)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def distributeCandies(self, n: int, limit: int) -> int:\n",
    "        \n",
    "#         @lru_cache(None)\n",
    "#         def dfs(i,j):\n",
    "#             if (3 - i) * limit < j:return 0\n",
    "#             if j == 0:return 1\n",
    "#             if i == 1:return max(0,min(limit,j) - max(0,(j - limit)) + 1)\n",
    "#             if i == 2:return j <= limit\n",
    "#             ans = 0\n",
    "#             for k in range(limit+1):\n",
    "#                 if k > j:break\n",
    "#                 ans += dfs(i+1,j-k)\n",
    "#             return ans\n",
    "        \n",
    "#         return dfs(0,n)\n",
    "    \n",
    "        dp = [[0 for _ in range(n+2)] for _ in range(5)]\n",
    "        for i in range(5):\n",
    "            dp[i][0] = 1\n",
    "        ans = 0\n",
    "        for j in range(max(0,n-limit),n+1):\n",
    "            if 2 * limit < j:\n",
    "                dp[1][j] = 0\n",
    "                break\n",
    "            # dp[1][j] = max(0,min(limit,j) - max(0,(j-limit)) + 1)\n",
    "            ans += max(0,min(limit,j) - max(0,(j-limit)) + 1)\n",
    "        return ans\n",
    "\n",
    "        # ans = 0\n",
    "        # if n <= limit:\n",
    "        #     for k in range(n+1):\n",
    "        #         ans += dp[1][n-k]\n",
    "        # else:\n",
    "        #     for k in range(limit+1):\n",
    "        #         ans += dp[1][n-k]\n",
    "        # return ans     "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def distributeCandies(self, n: int, limit: int) -> int:\n",
    "        ans = 0\n",
    "        for k in range(8):\n",
    "            t = bin(k).count('1')\n",
    "            rem = n - (limit + 1) * t\n",
    "            if rem < 0:\n",
    "                continue\n",
    "            ans += pow(-1, t) * comb(rem + 2, 2)\n",
    "        return ans\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "def c2(n: int) -> int:\n",
    "    return n*(n-1)//2 if n>1 else 0\n",
    "\n",
    "class Solution:\n",
    "    def distributeCandies(self, n: int, limit: int) -> int:\n",
    "        return c2(n+2)-3*c2(n-limit+1)+3*c2(n-2*limit)-c2(n-3*limit-1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def distributeCandies(self, n: int, limit: int) -> int:\n",
    "        if n > 3 * limit:\n",
    "            return 0\n",
    "        ans = comb(n + 2, 2)\n",
    "        if n > limit:\n",
    "            ans -= 3 * comb(n - limit + 1, 2)\n",
    "        if n - 2 >= 2 * limit:\n",
    "            ans += 3 * comb(n - 2 * limit, 2)\n",
    "        return ans\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def distributeCandies(self, n: int, limit: int) -> int:\n",
    "        if n > 3 * limit:\n",
    "            return 0\n",
    "        ans = comb(n + 2, 2)\n",
    "        if n > limit:\n",
    "            ans -= 3 * comb(n - limit + 1, 2)\n",
    "        if n - 2 >= 2 * limit:\n",
    "            ans += 3 * comb(n - 2 * limit, 2)\n",
    "        return ans\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def distributeCandies(self, n: int, limit: int) -> int:\n",
    "                return (math.comb(n + 2, 2) - (3 * math.comb(max(n - (limit + 1) + 2, 0), 2) \n",
    "                                       - 3 * math.comb(max(n - 2 * (limit + 1) + 2, 0), 2) \n",
    "                                       + math.comb(max(n - 3 * (limit + 1) + 2, 0), 2)))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "def c2(n: int) -> int:\n",
    "    return n * (n - 1) // 2 if n > 1 else 0\n",
    "\n",
    "class Solution:\n",
    "    def distributeCandies(self, n: int, limit: int) -> int:\n",
    "        return c2(n + 2) - 3 * c2(n - limit + 1) + 3 * c2(n - 2 * limit) - c2(n - 3 * limit - 1)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def distributeCandies(self, n: int, limit: int) -> int:\n",
    "        def cal(a, b):\n",
    "            if a < b:\n",
    "                return 0\n",
    "            return comb(a, b)\n",
    "        \n",
    "        return cal(n + 2, 2) - 3 * cal(n - limit + 1, 2) + 3 * cal(n - 2 * limit, 2) - cal(n - 3 * limit - 1, 2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "def c2(n: int) -> int:\n",
    "    return n * (n - 1) // 2 if n > 1 else 0\n",
    "\n",
    "class Solution:\n",
    "    def distributeCandies(self, n: int, limit: int) -> int:\n",
    "        return c2(n + 2) - 3 * c2(n - limit + 1) + 3 * c2(n - 2 * limit) - c2(n - 3 * limit - 1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def distributeCandies(self, n: int, limit: int) -> int:\n",
    "        def c(n):\n",
    "            if n < 2:\n",
    "                return 0\n",
    "            return comb(n, 2)\n",
    "        return comb(n + 2, 2) - 3 * c(n - (limit + 1) + 2) + 3 * c(n - 2 * (limit + 1) + 2) - c(n - 3 * (limit + 1) + 2)\n",
    "        \n",
    "        "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def distributeCandies(self, n: int, limit: int) -> int:\n",
    "        def c(n):\n",
    "            if n < 2:\n",
    "                return 0\n",
    "            return comb(n, 2)\n",
    "        return comb(n + 2, 2) - 3 * c(n - (limit + 1) + 2) + 3 * c(n - 2 * (limit + 1) + 2) - c(n - 3 * (limit + 1) + 2)\n",
    "        \n",
    "        "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "def C2(n: int) -> int:\n",
    "    if n <= 1:\n",
    "        return 0\n",
    "    return n * (n - 1) // 2\n",
    "\n",
    "\n",
    "class Solution:\n",
    "    def distributeCandies(self, n: int, limit: int) -> int:\n",
    "        res0 = C2(n + 2)  # 总方案数,x1+x2+x3=n的非负整数解的个数(隔板法)\n",
    "        res1 = 3 * C2(n + 2 - (limit + 1))  # 至少一个小朋友分到的糖果超过 limit 的方案数\n",
    "        res2 = 3 * C2(n + 2 - 2 * (limit + 1))  # 至少两个小朋友分到的糖果超过 limit\n",
    "        res3 = C2(n + 2 - 3 * (limit + 1))  # 三个小朋友分到的糖果都超过 limit\n",
    "        return res0 - res1 + res2 - res3\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "def c2(n: int) -> int:\n",
    "    return n * (n - 1) // 2 if n > 1 else 0\n",
    "\n",
    "class Solution:\n",
    "    def distributeCandies(self, n: int, limit: int) -> int:\n",
    "        return c2(n + 2) - 3 * c2(n - limit + 1) + 3 * c2(n - 2 * limit) - c2(n - 3 * limit - 1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def distributeCandies(self, n: int, limit: int) -> int:\n",
    "        ans = 0\n",
    "        for k in range(8):\n",
    "            t = k.bit_count()\n",
    "            rem = n - (limit + 1) * t\n",
    "            if rem >= 0:\n",
    "                ans += pow(-1, t) * comb(rem + 2, 2)\n",
    "        return ans\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def distributeCandies(self, n: int, limit: int) -> int:\n",
    "            \n",
    "        def helper(num):\n",
    "            if num < 0:\n",
    "                return 0\n",
    "            return (num + 2) * (num + 1) // 2\n",
    "\n",
    "        return helper(n) - (3 * helper(n - limit - 1) - 3 * helper(n - 2 * limit - 2) + helper(n - 3 * limit - 3))\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "# def c2(n: int) -> int:\n",
    "#     return n * (n - 1) // 2 if n > 1 else 0\n",
    "\n",
    "# class Solution:\n",
    "#     def distributeCandies(self, n: int, limit: int) -> int:\n",
    "#         print(c2(n + 2), c2(n - limit + 1), 3 * c2(n - 2 * limit), c2(n - 3 * limit - 1))\n",
    "#         return c2(n + 2) - 3 * c2(n - limit + 1) + 3 * c2(n - 2 * limit) - c2(n - 3 * limit - 1)\n",
    "\n",
    "\n",
    "import math\n",
    "class Solution:\n",
    "    def distributeCandies(self, n: int, lim: int) -> int:\n",
    "        \n",
    "        co = math.comb \n",
    "        # def co(n, y) -> int:\n",
    "        #     return n * (n - 1) // 2 if n > 1 else 0\n",
    "        all = co(n + 3 - 1, 2) \n",
    "        if n<=lim:\n",
    "            a = 0 \n",
    "        else:\n",
    "            res = (n - 1 - lim) \n",
    "            a = co(res + 3 - 1, 2)\n",
    "        b = a \n",
    "        c = a \n",
    "        if n<(2*lim + 2):\n",
    "            d = 0 \n",
    "        else:\n",
    "            res = n-2-2*lim \n",
    "            d = co(res + 3 - 1, 2) \n",
    "        e = d \n",
    "        f = d \n",
    "        if n<(3*lim + 3):\n",
    "            g = 0 \n",
    "        else:\n",
    "            res = n-3-3*lim \n",
    "            g = co(res+3-1, 2) \n",
    "        # print(all, a, d, g)\n",
    "        return all - (a+b+c-d-e-f+g)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def distributeCandies(self, n: int, limit: int) -> int:\n",
    "        if n > limit*3: return 0\n",
    "        hi = n-max(0,n-2*limit)\n",
    "        lo = max(0,n-limit)\n",
    "        def calc(x):\n",
    "            return min(limit,x)-max(0,x-limit)+1\n",
    "        def calc2(mi,mx):\n",
    "            return (mi+mx)*(mx-mi+1)//2\n",
    "        # print(lo,hi)\n",
    "        if hi <= limit:\n",
    "            return calc2(calc(lo),calc(hi))\n",
    "        elif lo >= limit:\n",
    "            return calc2(calc(hi),calc(lo))\n",
    "        else:\n",
    "            return calc2(calc(lo),calc(limit))+calc2(calc(hi),calc(limit))-limit-1\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "def c2(n: int) -> int:\n",
    "    return n * (n - 1) // 2 if n > 1 else 0\n",
    "\n",
    "class Solution:\n",
    "    def distributeCandies(self, n: int, limit: int) -> int:\n",
    "        return c2(n + 2) - 3 * c2(n - limit + 1) + 3 * c2(n - 2 * limit) - c2(n - 3 * limit - 1)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "def check(x):\n",
    "    if x <= 1:\n",
    "        return 0\n",
    "    return x*(x-1)//2\n",
    "\n",
    "class Solution:\n",
    "    def distributeCandies(self, n: int, limit: int) -> int:\n",
    "        return check(n+2) - 3*check(n-limit+1) + 3*check(n-2*limit) - check(n-3*limit-1)\n",
    "\n",
    "\n",
    "\n"
   ]
  }
 ],
 "metadata": {},
 "nbformat": 4,
 "nbformat_minor": 2
}
